import re
import sys
import os


class Struct2Proto:
    def __init__(self, proto_path, struct_path, with_hdr = False):
        self.proto_messages = {}
        self.struct_messages = {}
        self.proto_enums = {}
        self.struct_enums = set()
        self.proto_path = proto_path
        self.struct_path = struct_path
        self.with_hdr = with_hdr

    # Read the content of C struct file
    def ReadStructFile(self, file_path):
        # Read the content of C struct file
        with open(file_path, "r") as struct_fd:
            struct_file_content = struct_fd.read()
            # remove comments
            struct_file_content = re.sub(
                r"/\*.*?\*/", "", struct_file_content, flags=re.DOTALL
            )
            struct_file_content = re.sub(
                r"//.*?\n", "\n", struct_file_content
            )
        return struct_file_content
        
    # Read the content of proto file  
    def ReadProtoFile(self, file_path):
        # Read the content of proto file
        with open(file_path, "r") as proto_fd:
            proto_file_content = proto_fd.read()
            # remove the comments
            proto_file_content = re.sub(
                r"/\*.*?\*/", "", proto_file_content, flags=re.DOTALL
            )
            proto_file_content = re.sub(r"//.*?\n", "\n", proto_file_content)
        return proto_file_content
        
    def parse_struct_file(self, struct_file_content):
        struct_pattern = re.compile(
            r"""
                struct\s+         
                (?P<struct_name>[a-zA-Z_][a-zA-Z0-9_]*)\s*     # Capture the struct name 
                \{                 
                (?P<struct_fields>[^}]*)                     # Capture this struct's fields
                \};                
            """,
            re.DOTALL | re.VERBOSE,
        )
        field_pattern = re.compile(
            r"""
                (?P<field_type>[a-zA-Z_][a-zA-Z0-9_]*)\s+   # First group: field type
                (?P<field_name>[a-zA-Z_][a-zA-Z0-9_]*)\s*   # Second group: field name
                (?P<array_size>\[\d+u?\])?\s*;            # Optional third group: array size
            """,
            re.VERBOSE,
        )
        enum_pattern = re.compile(
            r"""
                enum\s+  
                (?P<enum_type>[a-zA-Z_][a-zA-Z0-9_]*)\s*     # Capture the enum type
                \{[^}]*\};
            """,
            re.VERBOSE | re.DOTALL,
        )

        for struct_match in struct_pattern.finditer(struct_file_content):
            struct_name = struct_match.group("struct_name")
            fields_content = struct_match.group("struct_fields")
            fields = []
            for match in field_pattern.finditer(fields_content):
                fields.append(
                    (
                        match["array_size"].strip("[]") if match["array_size"] else "",
                        match["field_type"],
                        match["field_name"],
                    ))
            self.struct_messages[struct_name] = fields

        self.struct_enums.update(set(enum_pattern.findall(struct_file_content)))

    def parse_typedef_file(self, struct_file_content):
        struct_pattern = re.compile(
            r"""
                typedef\s+struct\s+         
                (?P<struct_orignal_name>[a-zA-Z_][a-zA-Z0-9_]*)?\s*  # Capture the optional struct orignal name 
                \{                 
                (?P<struct_fields>[^}]*)                           # Capture this struct's fields
                \}\s*
                (?P<struct_name>[a-zA-Z_][a-zA-Z0-9_]*)\s*;          # Capture the struct name
            """,
            re.DOTALL | re.VERBOSE,
        )
        field_pattern = re.compile(
            r"""
                (?P<field_type>[a-zA-Z_][a-zA-Z0-9_]*)\s+   # First group: field type
                (?P<field_name>[a-zA-Z_][a-zA-Z0-9_]*)\s*   # Second group: field name
                (?P<array_size>\[\d+u?\])?\s*;            # Optional third group: array size
            """,
            re.VERBOSE,
        )
        enum_pattern = re.compile(
            r"""
                typedef\s+enum  
                (?P<enum_orignal_type>\s+[a-zA-Z_][a-zA-Z0-9_]*)?\s*   # Capture the optional enum_orignal_type
                \{[^}]*\}\s*
                (?P<enum_type>[a-zA-Z_][a-zA-Z0-9_]*)\s*             # Capture the enum type
            """,
            re.VERBOSE | re.DOTALL,
        )

        for struct_match in struct_pattern.finditer(struct_file_content):
            struct_name = struct_match.group("struct_name")
            print(f"struct name:{struct_name}\n\n")
            fields_content = struct_match.group("struct_fields")
            fields = []
            for match in field_pattern.finditer(fields_content):
                fields.append(
                    (
                        match["array_size"],
                        match["field_type"],
                        match["field_name"],
                    )
                )
            self.struct_messages[struct_name] = fields

        for match in enum_pattern.finditer(struct_file_content):
            self.struct_enums.add(match["enum_type"])

    def parse_proto_file(self, proto_file_content):
        message_pattern = re.compile(
            r"""
                message\s+
                (?P<message_name>[a-zA-Z_][a-zA-Z0-9_]*)\s*  # Capture the proto msg name
                {(?P<message_fields>[^}]*)}                # Capture the proto msg fields
            """,
            re.VERBOSE | re.DOTALL,
        )

        field_pattern = re.compile(
            r"""
                (?P<repeated>repeated\s+)?
                (?P<field_type>[a-zA-Z_][a-zA-Z0-9_]*)\s+
                (?P<field_name>[a-zA-Z_][a-zA-Z0-9_]*)\s*=\s*\d+;
            """,
            re.VERBOSE,
        )

        enum_pattern = re.compile(
            r"""
                enum\s+
                (?P<enum_name>[a-zA-Z_][a-zA-Z0-9_]*)\s*
                {[^}]*}
            """,
            re.VERBOSE,
        )
        package_pattern = re.compile(
            r"""
                package\s+
                (?P<package_name>[a-zA-Z0-9_.]+)\s*;
            """,
            re.VERBOSE,
        )

        package_name = package_pattern.findall(proto_file_content)
        # print(f"+++++++++{package_name}")
        if len(package_name) == 1:
            proto_namespace = package_name[0]
            # replace . in package name with ::
            proto_namespace = proto_namespace.replace(".", "::")
            proto_namespace += "::"
            # print(f"===== proto namespace = {proto_namespace}")
        else:
            print("Error: package name not found in proto file, or too many package names found")
            return False
        
        self.proto_messages[proto_namespace] = {}
        self.proto_enums[proto_namespace] = {}

        for message_match in message_pattern.finditer(proto_file_content):
            message_name = message_match.group("message_name")
            fields = field_pattern.findall(message_match.group("message_fields"))
            self.proto_messages[proto_namespace][message_name] = fields
        
        self.proto_enums[proto_namespace] = set(enum_pattern.findall(proto_file_content))

    def sanity_check(self, proto_messages, proto_enums, struct_messages, struct_enums):
        proto_diff_struct = proto_messages.keys() - struct_messages.keys()
        struct_diff_proto = struct_messages.keys() - proto_messages.keys()
        check_pass = True
        if proto_diff_struct:
            check_pass = False
            print(f"ProtoDiffStruct Error: {proto_diff_struct}")
        if struct_diff_proto:
            check_pass = False
            print(f"StructDiffProto Error: {struct_diff_proto}")

        if proto_enums != struct_enums:
            check_pass = False
            print(f"ProtoEnumDiff Error: {proto_enums} - {struct_enums}")

        if check_pass:
            for msg, proto_fields in proto_messages.items():
                struct_fields = struct_messages[msg]
                if len(proto_fields) == len(struct_fields):
                    tmp_proto_fields = sorted(proto_fields, key=lambda t: t[2])
                    tmp_struct_fields = sorted(struct_fields, key=lambda t: t[2])
                    for (proto_array, proto_type, proto_name), (
                        struct_array,
                        struct_type,
                        struct_name,
                    ) in zip(tmp_proto_fields, tmp_struct_fields):
                        if proto_name != struct_name:
                            check_pass = False
                            print(
                                f"Error: Mesasage field name is not matched, [{msg}:{proto_name}] in proto file, [{msg}:{struct_name}] in struct file"
                            )
                        #   if proto_type != struct_type:
                        #       check_pass = False
                        #       print(f"Error: Message field type is not matched!, [{msg}:{proto_name}:{proto_type}] in proto file, [{msg}:{struct_name}:{struct_type}] in struct file")
                        if (proto_array == "repeated") and (struct_array == ""):
                            check_pass = False
                            print(f"Error: Message filed array is not matched!")
                        if (proto_array == None) and (struct_array > ""):
                            check_pass = False
                            print(f"Error: Message filed array is not matched!")
                else:
                    check_pass = False
                    print(
                        f"Error: Message {msg}'s fields are not matched in proto file and struct file"
                    )

        return check_pass
    
    def generate_include(self):
        include = "#include <chrono>\n"
        if os.path.isdir(self.struct_path):
            for file in os.listdir(self.struct_path):
                if file.endswith(".h"):
                    include += f'#include "{file}"\n'
        elif self.struct_path.endswith(".h"):
            include += f'#include "{self.struct_path}"\n'
        
        if os.path.isdir(self.proto_path):
            for file in os.listdir(self.proto_path):
                if file.endswith(".proto"):
                    file.replace(".proto", "_idl_wrapper.h")
                    include += f'#include "{file}"\n'
        elif self.proto_path.endswith(".proto"):
            include += f'#include "{self.proto_path.replace(".proto", "_idl_wrapper.h")}"\n'
        return include
    
    def generate_common_function(self):
        code = ""
        code += "[[gnu::unused]] static uint64_t GetTimeStamp() {\n"
        code += "  auto time_now = std::chrono::system_clock::now();\n"
        code += "  auto time_ms = std::chrono::duration_cast<std::chrono::milliseconds>(\n"
        code += "      time_now.time_since_epoch());\n"
        code += "  return static_cast<uint64_t>(time_ms.count());\n"
        code += "}\n"
        return code

    def generate_cpp_conversion(self):
        basic_types = {
            "double",
            "float",
            "int32",
            "int64",
            "uint32",
            "uint64",
            "sint32",
            "sint64",
            "fixed32",
            "fixed64",
            "sfixed32",
            "sfixed64",
            "bool",
            "string",
            "bytes",
        }
        src_code = self.generate_include()
        if not self.with_hdr:
            src_code += "namespace {\n"
        src_code += self.generate_common_function()
        
        hd_code = ""
        if self.with_hdr:
            hd_code = self.generate_include()
        
        for namespace, messages in self.proto_messages.items():
            for msg, fields in messages.items():
                src_code += f"void Struct2Proto_{msg}(const {msg}& src, {namespace}{msg}* dst) {{\n"
                hd_code += f"void Struct2Proto_{msg}(const {msg}& src, {namespace}{msg}* dst);\n\n"
                for repeated, type, name in fields:
                    if repeated:
                        src_code += f"  for (const auto& struct_item : src.{name}) {{\n"
                        if type in basic_types:
                            src_code += f"      dst->add_{name.lower()}(struct_item);\n"
                        elif type in self.struct_enums:
                            src_code += f"      dst->set_{name.lower()}(static_cast<{namespace}{type}>(struct_item));\n"
                        else:
                            src_code += f"      Struct2Proto_{type}(struct_item, dst->add_{name.lower()}());\n"
                        src_code += "   }\n"
                    elif type in basic_types:
                        src_code += f"  dst->set_{name.lower()}(src.{name});\n"
                    elif type in self.struct_enums:
                        src_code += f"  dst->set_{name.lower()}(static_cast<{namespace}{type}>(src.{name}));\n"
                    else:
                        src_code += f"  Struct2Proto_{type}(src.{name}, dst->mutable_{name.lower()}());\n"
                src_code += "}\n\n"
        if not self.with_hdr:
            src_code += "}\n"
        return [src_code, hd_code]

    def run(self):
        if os.path.isdir(self.struct_path):
            for file in os.listdir(self.struct_path):
                if file.endswith(".h"):
                    struct_file_content = self.ReadStructFile(os.path.join(self.struct_path, file))
                    self.parse_struct_file(struct_file_content)
                    self.parse_typedef_file(struct_file_content)
        elif self.struct_path.endswith(".h"):
            struct_file_content = self.ReadStructFile(self.struct_path)
            self.parse_struct_file(struct_file_content)
            self.parse_typedef_file(struct_file_content)

        if os.path.isdir(self.proto_path):
            for file in os.listdir(self.proto_path):
                if file.endswith(".proto"):
                    proto_file_content = self.ReadProtoFile(os.path.join(self.proto_path, file))
                    self.parse_proto_file(proto_file_content)
        elif self.proto_path.endswith(".proto"):
            proto_file_content = self.ReadProtoFile(self.proto_path)
            self.parse_proto_file(proto_file_content)

        # print(f"strut enum: {self.struct_enums}")
        # print(f"proto enum: {self.proto_enums}")
        # print("\n\n")
        # print(f"struct: {self.struct_messages}")
        # print("\n\n")
        # print(f"proto: {self.proto_messages}")

        proto_messages = {}
        for _, msg in self.proto_messages.items():
            proto_messages.update(msg)
        proto_enums = set()
        for _,enums in self.proto_enums.items():
            proto_enums.update(enums)
        if self.sanity_check(proto_messages, proto_enums, self.struct_messages, self.struct_enums):
            return self.generate_cpp_conversion()
        else:
            print(
                "sanity check failed! please inspect the struct file and the proto file!"
            )

def main():
    # get input arguments by system argument
    input_proto = sys.argv[1]
    input_h = sys.argv[2]
    output_cpp = sys.argv[3]
    with_hdr = False
    if len(sys.argv) == 5:
        with_hdr = True
    generator = Struct2Proto(input_proto, input_h, with_hdr)
    [cpp_code, hd_code] = generator.run()
    if cpp_code:
        with open(output_cpp, "w") as fd:
            fd.write(cpp_code)
    else:
        print("Error: cpp code is empty!")
    if with_hdr and hd_code:
        output_h = output_cpp.replace(".cpp", ".h")
        with open(output_h, "w") as fd:
            fd.write(hd_code)


if __name__ == "__main__":
    main()
